import load
import numpy as np
import keras
import pickle

keras.initializers.RandomNormal(mean=0.0, stddev=0.05, seed=None)

class dingjiaModel:
    def __init__(self,model_file=None):
        model = keras.Sequential()
        model.add(keras.layers.Dense(1, input_shape=(2,), activation='relu'))
        model.compile(optimizer='adam', loss='mse')
        self.model = model

        if model_file:
            self.model = pickle.load(open(model_file, 'rb'))
            # self.model.set_weights(net_params)

    def train(self,allJingWei,allJiage):
        self.model.fit(allJingWei, allJiage, epochs=2500, batch_size=250)

    def predict(self,allJingWei):
        return self.model.predict(allJingWei)

    def save(self,name):
        pickle.dump(self.model, open(name+'.pkl', 'wb'), protocol=2)

allJingWeiA=[]
allJiageA=[]
allJingWeiB=[]
allJiageB=[]

for i in range(1,836):
    jingdu=float(load.getCellT1(i,load.jingdu))
    weidu=float(load.getCellT1(i,load.weidu))
    wancheng=int(load.getCellT1(i,load.wancheng))
    if wancheng==1:
        allJingWeiA.append([jingdu,weidu])
        allJiageA.append(float(load.getCellT1(i,load.jiage)))
    else:
        allJingWeiB.append([jingdu, weidu])
        allJiageB.append(float(load.getCellT1(i, load.jiage)))

allJingWeiA = np.array(allJingWeiA)
allJiageA = np.array(allJiageA)
allJingWeiB = np.array(allJingWeiB)
allJiageB = np.array(allJiageB)

if __name__=='__main__':
    modelA=dingjiaModel('model.pkl')
    modelB=dingjiaModel()
    modelB.train(allJingWeiB,allJiageB)
    print(modelA.model.get_weights())
    print(modelB.model.get_weights())
    modelB.save('model2')